/*
 * Mp4Muxer.cpp
 *
 *  Created on: 2021-08-02
 *      Author: chuanjiang.zh
 */

#include "Mp4Muxer.h"
#include "mov-format.h"
#include <errno.h>

static struct mov_buffer_t s_io = {
	[](void *ctx, void *data, uint64_t bytes) {
	MovMemoryBuffer *thiz = (MovMemoryBuffer *)ctx;
	return thiz->onRead(data, bytes);
	},
	[](void *ctx, const void *data, uint64_t bytes) {
		MovMemoryBuffer *thiz = (MovMemoryBuffer *)ctx;
		return thiz->onWrite(data, bytes);
	},
	[](void *ctx, uint64_t offset) {
		MovMemoryBuffer *thiz = (MovMemoryBuffer *)ctx;
		return thiz->onSeek(offset);
	},
	[](void *ctx) {
		MovMemoryBuffer *thiz = (MovMemoryBuffer *)ctx;
		return (uint64_t)thiz->onTell();
	}
};


Mp4Muxer::Mp4Muxer():
	m_writer(),
	m_videoTrack(-1),
	m_audioTrack(-1),
	m_keyWanted(true),
	m_lastTs()
{
}

Mp4Muxer::~Mp4Muxer()
{
	close();
}

int Mp4Muxer::open(const UMediaFormat& fmt)
{
	closeWriter();

	m_writer = fmp4_writer_create(&s_io, &m_buffer, MOV_FLAG_SEGMENT);
	if (!m_writer)
	{
		return EINVAL;
	}

	if (fmt.hasVideo())
	{
		uint8_t object = MOV_OBJECT_H264;
		if (fmt.m_codec == UMEDIA_CODEC_HEVC)
		{
			object = MOV_OBJECT_HEVC;
		}
		else if (fmt.m_codec == UMEDIA_CODEC_MPEG4)
		{
			object = MOV_OBJECT_MP4V;
		}
		else if (fmt.m_codec == UMEDIA_CODEC_VP8)
		{
			object = MOV_OBJECT_VP8;
		}

		m_videoTrack = fmp4_writer_add_video(m_writer, object, fmt.m_width, fmt.m_height, fmt.getVideoProp(), fmt.getVideoPropSize());
	}

	if (fmt.hasAudio())
	{
		uint8_t object = MOV_OBJECT_AAC;
		if (fmt.m_codec == UMEDIA_CODEC_AAC)
		{
			object = MOV_OBJECT_AAC;
		}
		else if (fmt.m_codec == UMEDIA_CODEC_G711A)
		{
			object = MOV_OBJECT_G711a;
		}
		else if (fmt.m_codec == UMEDIA_CODEC_G711A)
		{
			object = MOV_OBJECT_G711u;
		}
		else if (fmt.m_codec == UMEDIA_CODEC_OPUS)
		{
			object = MOV_OBJECT_OPUS;
		}

		m_audioTrack = fmp4_writer_add_audio(m_writer, object, 
			fmt.m_channels, fmt.m_sampleBits, fmt.m_sampleRate, fmt.getAudioConfig(), fmt.getAudioConfigSize());
	}

	fmp4_writer_init_segment(m_writer);

	m_keyWanted = true;

	return 0;
}

void Mp4Muxer::close()
{
	flush();

	closeWriter();
	
	m_keyWanted = true;
}

bool Mp4Muxer::isOpen()
{
	return (m_writer != nullptr);
}

bool Mp4Muxer::write(int type, uint8_t* data, int size, int64_t pts, int64_t dts, int flags)
{
	bool got = false;

	bool isKey = ((flags & UMEDIA_FLAG_KEY) == UMEDIA_FLAG_KEY);
	if (m_keyWanted)
	{
		if (type == UMEDIA_TYPE_VIDEO && isKey)
		{
			m_keyWanted = false;
		}
		else
		{
			return false;
		}
	}

	if (type == UMEDIA_TYPE_VIDEO)
	{
		if (isKey)
		{
			std::string data = m_buffer.getAndClearMemory();
			fireMuxerData(data, flags);
		}

		int movFlags = isKey ? MOV_AV_FLAG_KEYFREAME : 0;
		int ret = fmp4_writer_write(m_writer, m_videoTrack, data, size, pts, dts, movFlags);
		got = (ret == 0);

		m_lastTs = dts;
	}
	else if (type == UMEDIA_TYPE_AUDIO)
	{
		int movFlags = isKey ? MOV_AV_FLAG_KEYFREAME : 0;
		int ret = fmp4_writer_write(m_writer, m_audioTrack, data, size, pts, dts, movFlags);
		got = (ret == 0);
	}
	else
	{
		//
	}
	return got;
}

void Mp4Muxer::setSink(MuxerSink sink)
{
	m_sink = sink;
}

void Mp4Muxer::flush()
{
	std::string data = m_buffer.getAndClearMemory();
	fireMuxerData(data, 0);
}

void Mp4Muxer::closeWriter()
{
	if (m_writer)
	{
		fmp4_writer_destroy(m_writer);
		m_writer = nullptr;
	}

	m_videoTrack = -1;
	m_audioTrack = -1;
}

void Mp4Muxer::fireMuxerData(const std::string& data, int flags)
{
	if (m_sink)
	{
		m_sink((uint8_t*)data.c_str(), data.size(), m_lastTs, flags);
	}
}